﻿using System;
using System.Collections.Generic;
using System.IO;
using System.Runtime.InteropServices;
using ESRI.ArcGIS.Carto;
using ESRI.ArcGIS.DataSourcesGDB;
using ESRI.ArcGIS.DataSourcesRaster;
using ESRI.ArcGIS.esriSystem;
using ESRI.ArcGIS.GeoAnalyst;
using ESRI.ArcGIS.Geodatabase;
using ESRI.ArcGIS.Geometry;
using ESRI.ArcGIS.Geoprocessing;
using ESRI.ArcGIS.Geoprocessor;
using ESRI.ArcGIS.SpatialAnalystTools;

namespace Bes.Swsp.ToolBox
{
  public class ToolBox
  {
    private IGeometricNetwork _geometricNetwork;
    private IFeatureClass _catchmentClass;
    private IRaster _flowDirRaster;
    private IRaster _flowAccRaster;

    private StormDrainNetwork _stormNetwork;
    private IFeatureClass _nodesClass;
    private IFeatureClass _conveyanceClass;
    private bool _extractTerminalFlow = false;

    // Storm Drain Network Fields
    private const string ID_FIELD = "Id";
    private const string TYPE_FIELD = "Type";
    private const string DOWNNODE_FIELD = "DownNode";
    private const string UPNODE_FIELD = "UpNode";
    private const string UPNODES_FIELD = "UpNodes";
    private const string ENABLED_FIELD = "Enabled";
    private const string ERRORS_FIELD = "Errors";
    private const string SUBSURFACE_FIELD = "Subsurface";

    public event EventHandler<ProgressEventArgs> Progress;

    public IGeometricNetwork GeometricNetwork
    {
      get { return _geometricNetwork; }
      set { _geometricNetwork = value; }
    }

    public IFeatureClass CatchmentClass
    {
      get { return _catchmentClass; }
      set { _catchmentClass = value; }
    }

    public IRaster FlowDirRaster
    {
      get { return _flowDirRaster; }
      set { _flowDirRaster = value; }
    }

    public IRaster RunoffRaster
    {
      get { return _flowAccRaster; }
      set { _flowAccRaster = value; }
    }

    public IFeatureClass StormNetworkNodeClass
    {
      get { return _nodesClass; }
      set { _nodesClass = value; }
    }

    public IFeatureClass StormNetworkConveyanceClass
    {
      get { return _conveyanceClass; }
      set { _conveyanceClass = value; }
    }

    public bool ExtractTerminalFlow
    {
      get { return _extractTerminalFlow; }
      set { _extractTerminalFlow = value; }
    }

    public StormDrainNetwork StormNetwork
    {
      get { return _stormNetwork; }
      set { _stormNetwork = value; }
    }

    /// <summary>
    /// Accumulates values in input raster over surface and through underground conveyances, updating
    /// accumulation values to storm drain network elements.
    /// </summary>
    /// <param name="rasterToAccumulate"></param>
    public void AccumulateRaster(IRaster rasterToAccumulate)
    {
      // Accumulate consequences that "flow through" underground pipes
      IRaster localAccumulation = _CalculateLocalFlowAccumulation(rasterToAccumulate);
      _ExtractLocalAccumulationArc(_stormNetwork.Catchments, localAccumulation);
      _stormNetwork.Accumulate();

      // Add local accumulation for surface nodes
      _AddLocalAccumulation(_stormNetwork.Nodes, localAccumulation);
    }

    /// <summary>
    /// Extracts accumulation of values to the storm drain network, passing local accumulation
    /// through underground conveyances, updating accumulation values in storm drain network elements.
    /// </summary>
    /// <param name="localAccumulation">Raster representing the surface accumulation of a variable
    /// through the network.</param>
    /// <remarks>This method assumes the local accumulation and the flow direction grid were generated
    /// (or are in the format produced by) TauDEM. This means the flow direction uses direction codes
    /// of 1-8 and NODATA for inlets. Local accumulation would then have NODATA at inlets as well. This
    /// method will calculate total accumulation flowing into inlets by checking surrounding cells and
    /// adding accumulation from those flowing into the inlet. The inlet cell itself does not contribute
    /// anything to the accumulation.</remarks>
    public void AccumulateThroughNetwork(IRaster localAccumulation)
    {
      _ExtractLocalAccumulationTau(_stormNetwork.Catchments, localAccumulation);
      _stormNetwork.Accumulate();

      // Add local accumulation for surface nodes
      _AddLocalAccumulation(_stormNetwork.Nodes, localAccumulation);
    }

    /// <summary>
    /// Recalculates network flows and loads using given decumulation parameters
    /// </summary>
    /// <param name="nodeIds">The node Ids with volume reduction and/or effluent concentration.</param>
    /// <param name="volumeReductionPercent">The volume reduction for each node in the order of the given nodeIds.</param>
    /// <param name="effluentConcentrations">The effluent concentration for each node in the order of the give nodeIds.</param>
    public void Decumulate(int[] nodeIds, double[] volumeReductionPercent, double[] effluentConcentrations)
    {
      if (volumeReductionPercent.Length != nodeIds.Length || effluentConcentrations.Length != nodeIds.Length)
      {
        throw new ArgumentException("Length of decumulator parameter arrays are not equal.");
      }

      // Reset decumulation parameters
      foreach (NodePoint node in _stormNetwork.Nodes)
      {
        node.VolumeReduction = 0.0;
        node.ResetEffluentConcentration();
      }

      for (int i = 0; i < nodeIds.Length; i++)
      {
        NodePoint node = _stormNetwork.Nodes[nodeIds[i]];

        node.VolumeReduction = volumeReductionPercent[i] / 100.0;
        if(!double.IsNaN(effluentConcentrations[i]))
          node.EffluentConcentration = effluentConcentrations[i];
      }

      _stormNetwork.Accumulate(false);
    }

    public void UpdateEids()
    {
      bool cancel = false;

      int total = 0;
      int progress = 0;

      OnProgress(new ProgressEventArgs(0, 1, "Updating EIDs...", true), true);

      IEnumFeatureClass junctionClasses = _geometricNetwork.get_ClassesByType(esriFeatureType.esriFTSimpleJunction);
      IFeatureClass junctionClass = junctionClasses.Next();
      while (junctionClass != null)
      {
        int eidField = junctionClass.FindField("EID");
        if (eidField != -1)
        {
          total += junctionClass.FeatureCount(null);
        }
        ComUtility.ReleaseComObject(junctionClass);
        junctionClass = junctionClasses.Next();
      }

      int increment = _CalcIncrement(total);

      junctionClasses.Reset();
      junctionClass = junctionClasses.Next();
      while (junctionClass != null)
      {
        IWorkspaceEdit workspaceEdit = ((IDataset)junctionClass).Workspace as IWorkspaceEdit;
        workspaceEdit.StartEditing(false);
        IFeatureCursor cursor = null;
        try
        {
          int eidField = junctionClass.FindField("EID");
          if (eidField != -1)
          {
            workspaceEdit.StartEditOperation();
            try
            {
              cursor = junctionClass.Update(null, false);
              IFeature node = cursor.NextFeature();
              while (node != null)
              {
                try
                {
                  ISimpleJunctionFeature junction = node as ISimpleJunctionFeature;

                  object dbValue = node.get_Value(eidField);
                  if (dbValue != DBNull.Value)
                  {
                    int prevEID = Convert.ToInt32(dbValue);
                    if (prevEID != junction.EID)
                    {
                      Logger.Log(string.Format("Changing EID {0} to {1}.", dbValue, junction.EID));
                    }
                  }

                  node.set_Value(eidField, junction.EID);
                  cursor.UpdateFeature(node);

                  if (++progress % increment == 0)
                  {
                    ProgressEventArgs eventArgs = new ProgressEventArgs(progress, total, "Updating EIDs...", true);
                    OnProgress(eventArgs);
                    if (eventArgs.Cancel)
                      cancel = true;
                  }
                }
                finally
                {
                  ComUtility.ReleaseComObject(node);
                }

                if (cancel)
                  break;

                node = cursor.NextFeature();
              }
              workspaceEdit.StopEditOperation();
            }
            catch
            {
              workspaceEdit.AbortEditOperation();
              throw;
            }
            workspaceEdit.StopEditing(!cancel);
          }
        }
        catch
        {
          workspaceEdit.StopEditing(false);
          throw;
        }
        finally
        {
          ComUtility.ReleaseComObject(cursor);
        }

        ComUtility.ReleaseComObject(junctionClass);
        junctionClass = junctionClasses.Next();
      }

      OnProgress(cancel
        ? new ProgressEventArgs(0, 0, "Canceled. Changes rolled back.")
        : new ProgressEventArgs(1, 1, string.Format("{0} EIDs updated.", total)), true);
    }

    [Obsolete("Now done internally by the StormDrainNetwork initialization.")]
    public void MarkDownstreamEnds()
    {
      INetwork network = _geometricNetwork.Network;
      INetTopology netTopology = network as INetTopology;

      int endsCount = 0;
      int total = 0;
      int progress = 0;

      OnProgress(new ProgressEventArgs(0, 1, "Processing nodes..."), true);

      IEnumFeatureClass junctionClasses = _geometricNetwork.get_ClassesByType(esriFeatureType.esriFTSimpleJunction);
      IFeatureClass junctionClass = junctionClasses.Next();
      while (junctionClass != null)
      {
        int dsFlagField = junctionClass.FindField("dsflag");
        if (dsFlagField != -1)
        {
          total += junctionClass.FeatureCount(null);
        }
        ComUtility.ReleaseComObject(junctionClass);
        junctionClass = junctionClasses.Next();
      }

      int increment = _CalcIncrement(total);

      junctionClasses.Reset();
      junctionClass = junctionClasses.Next();
      while (junctionClass != null)
      {
        IWorkspaceEdit workspaceEdit = ((IDataset)junctionClass).Workspace as IWorkspaceEdit;
        workspaceEdit.StartEditing(false);
        IFeatureCursor cursor = null;
        try
        {
          int edgeId;
          bool towardJunction;
          int dsFlagField = junctionClass.FindField("dsflag");

          if (dsFlagField != -1)
          {
            workspaceEdit.StartEditOperation();
            try
            {
              cursor = junctionClass.Update(null, false);
              IFeature node = cursor.NextFeature();
              while (node != null)
              {
                try
                {
                  ISimpleJunctionFeature junction = node as ISimpleJunctionFeature;
                  bool isDownstream = junction.EdgeFeatureCount > 0;
                  for (int i = 0; i < junction.EdgeFeatureCount; i++)
                  {
                    netTopology.GetAdjacentEdge(junction.EID, i, out edgeId, out towardJunction);
                    if (!towardJunction)
                    {
                      isDownstream = false;
                      break;
                    }
                  }

                  if (isDownstream)
                  {
                    endsCount++;
                    node.set_Value(dsFlagField, 1);
                    cursor.UpdateFeature(node);
                  }

                  if (++progress % increment == 0)
                  {
                    OnProgress(new ProgressEventArgs(progress, total, "Processing nodes..."));
                  }
                }
                finally
                {
                  ComUtility.ReleaseComObject(node);
                }
                node = cursor.NextFeature();
              }
              workspaceEdit.StopEditOperation();
            }
            catch
            {
              workspaceEdit.AbortEditOperation();
              throw;
            }
            workspaceEdit.StopEditing(true);
          }
        }
        catch
        {
          workspaceEdit.StopEditing(false);
          throw;
        }
        finally
        {
          ComUtility.ReleaseComObject(cursor);
        }

        ComUtility.ReleaseComObject(junctionClass);
        junctionClass = junctionClasses.Next();
      }

      OnProgress(new ProgressEventArgs(1, 1, string.Format("{0} downstream ends found.", endsCount)), true);
    }

    public void FindFlippedInletLeads()
    {
      OnProgress(new ProgressEventArgs(0, 1, "Finding flipped inlet leads...", true), true);

      IFeatureClass dmeClass = null;
      IEnumDataset datasets = _geometricNetwork.FeatureDataset.Subsets;
      IDataset dataset = datasets.Next();
      while (dataset != null)
      {
        if (dataset.Name == "dme_lines")
        {
          dmeClass = dataset as IFeatureClass;
          break;
        }
        ComUtility.ReleaseComObject(dataset);
        dataset = datasets.Next();
      }

      int progress = 0;
      int total = dmeClass.FeatureCount(null);
      int increment = _CalcIncrement(total);

      INetwork network = _geometricNetwork.Network;

      const string FLIPPED_FIELD = "Flipped";
      int flipCount = 0;
      int flipField = dmeClass.FindField(FLIPPED_FIELD);
      if (flipField == -1)
      {
        FeatureClassUtility.AddField(ref dmeClass, FLIPPED_FIELD, esriFieldType.esriFieldTypeSmallInteger, false);
        flipField = dmeClass.FindField(FLIPPED_FIELD);
      }

      List<string> inletSymbols = new List<string>(new string[] { "INLETS", "TRASH RACK", "DISCHARGE POINTS" });

      bool cancel = false;
      IWorkspaceEdit workspaceEdit = ((IDataset)dmeClass).Workspace as IWorkspaceEdit;
      workspaceEdit.StartEditing(false);
      try
      {
        workspaceEdit.StartEditOperation();
        try
        {
          IFeatureCursor cursor = null;
          try
          {
            cursor = dmeClass.Update(null, false);

            IFeature inletLead = cursor.NextFeature();
            while (inletLead != null)
            {
              IEdgeFeature edge = inletLead as IEdgeFeature;

              bool fromInlet = false;
              bool toInlet = false;
              bool flipped = false;

              IFeature fromNode = edge.FromJunctionFeature as IFeature;
              if (((ISimpleJunctionFeature)fromNode).EdgeFeatureCount == 1)
              {
                fromInlet = true;
              }
              //else
              //{
              //    int fromSymbolField = fromNode.Fields.FindField("SYMBOL_GROUP");
              //    if (fromSymbolField != -1)
              //    {
              //        string fromSymbol = fromNode.get_Value(fromSymbolField).ToString();
              //        if (inletSymbols.Contains(fromSymbol))
              //        {
              //            fromInlet = true;
              //        }
              //    }
              //}
              ComUtility.ReleaseComObject(fromNode);

              IFeature toNode = edge.ToJunctionFeature as IFeature;
              if (((ISimpleJunctionFeature)toNode).EdgeFeatureCount == 1)
              {
                //toInlet = true;
                int toSymbolField = toNode.Fields.FindField("SYMBOL_GROUP");
                if (toSymbolField != -1)
                {
                  string toSymbol = toNode.get_Value(toSymbolField).ToString();
                  if (inletSymbols.Contains(toSymbol))
                  {
                    toInlet = true;
                  }
                }
              }
              ComUtility.ReleaseComObject(toNode);

              if (toInlet && !fromInlet)
              {
                flipped = true;
                flipCount++;
              }

              inletLead.set_Value(flipField, flipped ? (short)1 : (short)0);

              try
              {
                cursor.UpdateFeature(inletLead);
              }
              catch
              {
                Logger.Log(string.Format("ToolBox.FindFlippedInletLeads: Error updating OID {0}.", inletLead.OID));
                throw;
              }

              if (++progress % increment == 0)
              {
                ProgressEventArgs eventArgs = new ProgressEventArgs(progress, total, string.Format("Finding flipped inlet leads... {0} found.", flipCount), true);
                OnProgress(eventArgs);
                if (eventArgs.Cancel)
                {
                  cancel = true;
                  break;
                }
              }

              inletLead = cursor.NextFeature();
            }
            cursor.Flush();
          }
          finally
          {
            ComUtility.ReleaseComObject(cursor);
          }
          workspaceEdit.StopEditOperation();
        }
        catch
        {
          workspaceEdit.AbortEditOperation();
          throw;
        }
        workspaceEdit.StopEditing(!cancel);
      }
      catch
      {
        workspaceEdit.StopEditing(false);
        throw;
      }

      OnProgress(cancel
        ? new ProgressEventArgs(0, 0, "Canceled. Changes rolled back.")
        : new ProgressEventArgs(1, 1, string.Format("{0} flipped inlet leads found.", flipCount)), true);
    }

    private IRaster _CalculateLocalFlowAccumulation(IRaster rasterToAccumulate)
    {
      #warning Redundant with StormDrainNetworkLoader._CalculateLocalFlowAccumulation
      OnProgress(new ProgressEventArgs(0, 1, "Calculating local accumulation..."), true);

      Geoprocessor gp = GeoprocessingUtility.GetGeoprocessor(true, false, true, _flowDirRaster);
      try
      {
        FlowAccumulation flowAccTool = new FlowAccumulation();
        flowAccTool.in_flow_direction_raster = _flowDirRaster;
        if(rasterToAccumulate != null)
          flowAccTool.in_weight_raster = rasterToAccumulate;
        flowAccTool.out_accumulation_raster = "local_accumulation";

        object result = GeoprocessingUtility.RunGpTool(gp, flowAccTool, OnProgress);
        Logger.Log(string.Format("Local flow accumulation: {0}", result));

        OnProgress(new ProgressEventArgs(0, 1, "Returning the local accumulation raster..."), true);

        string workspacePath = System.IO.Path.GetDirectoryName((string)result);
        string rasterName = System.IO.Path.GetFileName((string)result);

        IWorkspaceFactory wsf = new FileGDBWorkspaceFactoryClass();
        IWorkspace ws = wsf.OpenFromFile(workspacePath, 0);
        IRasterDataset rasterDataset = ((IRasterWorkspaceEx)ws).OpenRasterDataset(rasterName);

        return rasterDataset.CreateDefaultRaster();
      }
      catch
      {
        for (int i = 0; i < gp.MessageCount; i++)
          OnProgress(new ProgressEventArgs(0, 1, gp.GetMessage(i)), true);

        throw;
      }
      finally
      {
        GeoprocessingUtility.ResetGeoprocessor();
      }
    }

    private void _ExtractLocalAccumulationArc(ICollection<Catchment> catchments, IRaster localAccumulationRaster)
    {
#warning Redundant with StormDrainNetworkLoader._ExtractLocalAccumulation(). Consider merging.
      OnProgress(new ProgressEventArgs(0, 1, "Extracting local flow accumulation...", true), true);

      int progress = 0;
      int total = catchments.Count;
      int increment = _CalcIncrement(total);

      // order for performance (not a huge gain, unfortunately...)
      List<Catchment> orderedCatchments = new List<Catchment>(catchments);
      orderedCatchments.Sort((c1, c2) => c2.DrainageNode.Point.Y.CompareTo(c1.DrainageNode.Point.Y));

      foreach (Catchment catchment in orderedCatchments)
      {
        double accumulation = RasterUtility.GetPixelValueDouble(localAccumulationRaster, catchment.DrainageNode.Point);
        catchment.LocalFlowAccumulation = double.IsNaN(accumulation) ? 0.0 : accumulation;

        if (++progress % increment == 0)
        {
          ProgressEventArgs eventArgs = new ProgressEventArgs(progress, total, "Extracting local flow accumulation...", true);
          OnProgress(eventArgs);
          if (eventArgs.Cancel)
            throw new UserCanceledException();
        }
      }

      OnProgress(new ProgressEventArgs(1, 1, "Local flow accumulation extracted."), true);
    }

    private void _ExtractLocalAccumulationTau(ICollection<Catchment> catchments, IRaster localAccumulationRaster)
    {
      #warning Redundant with StormDrainNetworkLoader._ExtractLocalAccumulation(). Consider merging.
      OnProgress(new ProgressEventArgs(0, 1, "Extracting local flow accumulation...", true), true);

      int progress = 0;
      int total = catchments.Count;
      int increment = _CalcIncrement(total);

      // order for performance (not a huge gain, unfortunately...)
      List<Catchment> orderedCatchments = new List<Catchment>(catchments);
      orderedCatchments.Sort((c1, c2) => c2.DrainageNode.Point.Y.CompareTo(c1.DrainageNode.Point.Y));

      foreach (Catchment catchment in orderedCatchments)
      {
        catchment.LocalFlowAccumulation = ExtractTauAccumulationAtNullCell(localAccumulationRaster, _flowDirRaster, catchment.DrainageNode.Point);

        if (++progress % increment == 0)
        {
          ProgressEventArgs eventArgs = new ProgressEventArgs(progress, total, "Extracting local flow accumulation...", true);
          OnProgress(eventArgs);
          if (eventArgs.Cancel)
            throw new UserCanceledException();
        }
      }

      OnProgress(new ProgressEventArgs(1, 1, "Local flow accumulation extracted."), true);
    }

    /// <summary>
    /// This method assumes accumulation raster is generated by TauDEM i.e. accumlation includes current cell
    /// ArcGIS calculates accumulation without current cell
    /// </summary>
    /// <param name="accumRaster">Flow Accumulation grid created by TauDEM. (Accumulation values include
    /// accumulation from current cell.</param>
    /// <param name="flowdir">Flow Direction grid using TauDEM direction codes (i.e. 1-8)</param>
    /// <param name="location">The location of the point for which to get flow accumulation values. Assumes
    /// this point is at a NO DATA cell with flow entering from surrounding cells. (e.g. an inlet)</param>
    /// <returns>The flow accumulation for the cell at the specified location. Includes accumulated values
    /// flowing into the current cell, but not the accumulation at the current cell.</returns>
    public static double ExtractTauAccumulationAtNullCell(IRaster accumRaster, IRaster flowdir, IPoint location)
    {
      double cellWidth = ((IRasterAnalysisProps)accumRaster).PixelWidth;
      double cellHeight = ((IRasterAnalysisProps)accumRaster).PixelHeight;

      int[] xAdjust = new int[] { 0, 1, 1, 0, -1, -1, -1, 0, 1 };
      int[] yAdjust = new int[] { 0, 0, -1, -1, -1, 0, 1, 1, 1 };
      int[] reverseD = new int[] { 0, 5, 4, 3, 2, 1, 8, 7, 6 };

      double accumulation = 0.0;
      for (int d = 1; d < 9; d++)
      {
        IPoint dp = new PointClass();
        dp.X = location.X + (cellWidth * xAdjust[d]);
        dp.Y = location.Y + (cellHeight * yAdjust[d]);

        int direction = RasterUtility.GetPixelValueInt(flowdir, dp);
        if (direction == reverseD[d])
        {
          double pixelValue = RasterUtility.GetPixelValueDouble(accumRaster, dp);
          if (!double.IsNaN(pixelValue))
            accumulation += pixelValue;
        }
      }

      return accumulation;
    }

    private void _AddLocalAccumulation(ElementCollection<NodePoint> nodes, IRaster localAccumulationRaster, bool includeInlets = false)
    {
      // This method assumes inlets already have local flow accumulation set in the catchment. To add this to the inlet's
      // total accumulation set includeInlets to true. Make sure this is not being double counted.

      OnProgress(new ProgressEventArgs(0, 1, "Adding local accumulation...", true), true);

      int progress = 0;
      int total = nodes.Count;
      int increment = _CalcIncrement(total);

      // order for performance (not a huge gain, unfortunately...)
      List<NodePoint> orderedNodes = new List<NodePoint>(nodes);
      orderedNodes.Sort((n1, n2) => n1.Point.Y.CompareTo(n2.Point.Y));

      foreach (NodePoint node in orderedNodes)
      {
        // Inlet nodes are assumed to already be extracted in the catchment area
        if (node is InletNode)
        {
          if (includeInlets)
            node.TotalAccumulatedFlow += ((InletNode)node).CatchmentArea.LocalFlowAccumulation;
        }
        else if (node.HasSurfaceInflow)
        {
          double pixelValue = RasterUtility.GetPixelValueDouble(localAccumulationRaster, node.Point);
          node.TotalAccumulatedFlow += double.IsNaN(pixelValue) ? 0.0 : pixelValue;
        }

        if (++progress % increment == 0)
        {
          ProgressEventArgs eventArgs = new ProgressEventArgs(progress, total, "Adding local accumulation...", true);
          OnProgress(eventArgs);
          if (eventArgs.Cancel)
            throw new UserCanceledException();
        }
      }

      OnProgress(new ProgressEventArgs(1, 1, "Local accumulation added."), true);
    }

    internal void TransferReceivingAttribute(IFeatureClass watersClass, string watersAttribute, bool traceToUltimate)
    {
      System.Diagnostics.Stopwatch sw = new System.Diagnostics.Stopwatch();
      sw.Start();

      OnProgress(new ProgressEventArgs(0, 1, "Detecting nodes in receiving waters...", true), true);

      // Clear tags
      foreach (NodePoint node in _stormNetwork.Nodes)
        node.Tag = null;
      foreach (Conveyance conveyance in _stormNetwork.Conveyances)
        conveyance.Tag = null;

      int attributeIndex = watersClass.FindField(watersAttribute);

      // Get list of possible end-nodes
      List<NodePoint> nodeList = new List<NodePoint>();
      foreach (NodePoint node in _stormNetwork.Nodes)
      {
        // NOTE: Disabled nodes are allowed so that UICs can be included as end points.

        // Ignore upstream ends
        if (node.UpstreamConveyances.Count == 0)
          continue;

        // If node is completely underground, ignore (unless it is a UIC--i.e. disabled)
        if (node.Enabled)
        {
          if (node.DownstreamConveyance != null && node.DownstreamConveyance.IsSubSurface)
          {
            bool allSubsurface = true;
            foreach (Conveyance upConveyance in node.UpstreamConveyances)
            {
              if (!upConveyance.IsSubSurface)
              {
                allSubsurface = false;
                break;
              }
            }
            if (allSubsurface)
              continue;
          }
        }

        nodeList.Add(node);
      }

      int progress = 0;
      int total = nodeList.Count;
      int increment = _CalcIncrement(total);

      System.Diagnostics.Debug.WriteLine("Node list: " + sw.Elapsed.TotalSeconds.ToString());

      // Index waters polygons
      IFeatureIndex indexer = new FeatureIndexClass();
      indexer.FeatureClass = watersClass;
      indexer.Index(null, ((IGeoDataset)watersClass).Extent);
      IIndexQuery2 indexQuery = (IIndexQuery2)indexer;

      System.Diagnostics.Debug.WriteLine("Indexing: " + sw.Elapsed.TotalSeconds.ToString());

      // Collect waters attributes
      Dictionary<int, object> attributeDictionary = new Dictionary<int, object>();

      IFeatureCursor cursor = watersClass.Search(null, false);
      try
      {
        IFeature water = cursor.NextFeature();
        while (water != null)
        {
          try
          {
            attributeDictionary.Add(water.OID, water.get_Value(attributeIndex));
          }
          finally
          {
            ComUtility.ReleaseComObject(water);
          }
          water = cursor.NextFeature();
        }
      }
      finally
      {
        ComUtility.ReleaseComObject(cursor);
      }

      System.Diagnostics.Debug.WriteLine("Read attributes: " + sw.Elapsed.TotalSeconds.ToString());

      Dictionary<NodePoint, int> waterNodes = new Dictionary<NodePoint, int>();

      // Find nodes at receiving waters
      foreach (NodePoint node in nodeList)
      {
        int waterOid = indexQuery.WithinFeature(node.Point);
        if (waterOid != -1)
        {
          foreach (Conveyance upConveyance in node.UpstreamConveyances)
          {
            // Don't trace up paths that are completely underground (except in the case of UIC's--i.e. disabled nodes)
            if (node.Enabled && node.DownstreamConveyance != null && node.DownstreamConveyance.IsSubSurface && upConveyance.IsSubSurface)
              continue;

            waterNodes.Add(node, waterOid);
            break;
          }
        }

        if (++progress % increment == 0)
        {
          ProgressEventArgs eventArgs = new ProgressEventArgs(progress, total, "Detecting nodes in receiving waters...", true);
          OnProgress(eventArgs);
          if (eventArgs.Cancel)
            throw new UserCanceledException();
        }
      }

      OnProgress(new ProgressEventArgs(1, 1, string.Format("{0} receiving water nodes detected.", waterNodes.Count), true));

      System.Diagnostics.Debug.WriteLine("Detecting nodes: " + sw.Elapsed.TotalSeconds.ToString());

      OnProgress(new ProgressEventArgs(0, 1, "Transferring attributes...", true), true);

      progress = 0;
      total = waterNodes.Count;
      increment = _CalcIncrement(total);

      // Transfer attributes
      foreach (NodePoint node in waterNodes.Keys)
      {
        int waterOid = waterNodes[node];

        if (node.Tag != null)
          continue;

        foreach (Conveyance upConveyance in node.UpstreamConveyances)
        {
          // Need to recheck for underground conveyances
          if (node.Enabled && node.DownstreamConveyance != null && node.DownstreamConveyance.IsSubSurface && upConveyance.IsSubSurface)
            continue;

          _stormNetwork.ResetTraversedNodes();

          object attributeValue = attributeDictionary[waterOid];
          node.Tag = attributeValue;
          node.Traversed = true;

          _TransferAttributeUpstream(upConveyance, attributeValue, waterNodes, traceToUltimate);
        }

        if (++progress % increment == 0)
        {
          ProgressEventArgs eventArgs = new ProgressEventArgs(progress, total, "Transferring attributes...", true);
          OnProgress(eventArgs);
          if (eventArgs.Cancel)
            throw new UserCanceledException();
        }
      }

      OnProgress(new ProgressEventArgs(1, 1, "Receiving water attributes transferred."), true);

      sw.Stop();
      System.Diagnostics.Debug.WriteLine("Transfer Attribute:" + sw.Elapsed.TotalSeconds.ToString());
    }

    private void _TransferAttributeUpstream(Conveyance conveyance, object attributeValue, Dictionary<NodePoint, int> waterNodes, bool traceToUltimate)
    {
      if (!conveyance.Enabled)
        return;

      conveyance.Tag = attributeValue;

      // If we are only tracing to inital waters, then stop if upstream node is within another water body
      if (!traceToUltimate && waterNodes.ContainsKey(conveyance.UpstreamNode))
        return;

      conveyance.UpstreamNode.Tag = attributeValue;

      // Break out of loops / stop at disabled nodes
      if (conveyance.UpstreamNode.Traversed || !conveyance.UpstreamNode.Enabled)
        return;

      conveyance.UpstreamNode.Traversed = true;
      foreach (Conveyance upConveyance in conveyance.UpstreamNode.UpstreamConveyances)
      {
        _TransferAttributeUpstream(upConveyance, attributeValue, waterNodes, traceToUltimate);
      }
    }

    private int _CalcIncrement(int total)
    {
      int increment = total / 1000;
      if (increment == 0)
        increment = 1;
      return increment;
    }

    /// <summary>
    /// Sets the output environment for this raster tool, using the specified extent for output.
    /// </summary>
    /// <param name="pRasterTool">The tool to set the ouput environment for.</param>
    /// <param name="extentProvider">An IEnvelope or RasterDataset that defines the output extent.</param>
    /// <remarks>The extentProvider determines the extent, but the default raster (the input flow direction grid) determines cell size
    /// and snap registration of the cells.</remarks>
    internal void SetAnalysisEnvironment(IRasterAnalysisEnvironment rasterTool, object extentProvider)
    {
      object inputGrid = _flowDirRaster;

      rasterTool.DefaultOutputRasterPrefix = "StrmTlsTmp";
      rasterTool.OutWorkspace = GetTempRasterWorkspace();
      rasterTool.OutSpatialReference = ((IGeoDataset)inputGrid).SpatialReference;
      rasterTool.SetCellSize(esriRasterEnvSettingEnum.esriRasterEnvValue, ref inputGrid);
      rasterTool.SetExtent(esriRasterEnvSettingEnum.esriRasterEnvValue, ref extentProvider, ref inputGrid);
    }

    internal IWorkspace GetTempRasterWorkspace()
    {
      IWorkspaceFactory pWSFactory = new RasterWorkspaceFactoryClass();
      return pWSFactory.OpenFromFile(_GetNewTempDir(), 0);
    }

    private string _GetNewTempDir()
    {
      string documents = Environment.GetFolderPath(Environment.SpecialFolder.MyDocuments);
      string tempDir = System.IO.Path.Combine(documents, "temp");
      tempDir = System.IO.Path.Combine(tempDir, Guid.NewGuid().ToString("D"));
      Directory.CreateDirectory(tempDir);
      return tempDir;
    }

    internal static string CreateTempFileName(string directoryPath, string baseName, string extension)
    {
      DirectoryInfo directory = new DirectoryInfo(directoryPath);
      if (!directory.Exists)
      {
        throw new DirectoryNotFoundException();
      }

      if (extension == null)
      {
        extension = string.Empty;
      }

      if (extension.Length > 0 && !extension.StartsWith("."))
      {
        extension = "." + extension;
      }

      int subItemCount = directory.GetFileSystemInfos(baseName + "*").Length;

      string result = null;
      string fileName = null;
      for (int i = 1; i < subItemCount + 2; i++)
      {
        fileName = baseName + i.ToString(System.Globalization.CultureInfo.CurrentCulture) + extension;
        result = System.IO.Path.Combine(directoryPath, fileName);
        if (!File.Exists(result) && !Directory.Exists(result))
          return result;
      }

      throw new Exception("Could not create temporary file name!");
    }

    protected virtual void OnProgress(ProgressEventArgs eventArgs, bool log = false)
    {
      if (Progress != null)
      {
        Progress(this, eventArgs);
      }

      if (log)
      {
        Logger.Log(eventArgs.Message);
      }
    }
  }
}